package tv.geir.slackspace.code;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

import tv.geir.slackspace.common.BeanFactory;


/**
 * When hooked into the JVM this class captures bytecode-definition for all
 * classes that are loaded. 
 * <p>
 * The definitions are stored in <code>classDefinitions</code> and can be
 * accessed with <code>BytecodeAgent.getInstance().getBytecode( "classname" );</code>
 * where 'classname' is the canonical name for the class.
 * <p>
 * For BytecodeAgent to be hooked into the JVM the following argument must be passed
 * to JVM at startup:
 * 
 * <pre>-javaagent:"(..)\slackspaces-VERSION.jar"</pre> where VERSION corresponds
 * to the SlackSpaces version you are using.
 * 
 * 
 * @author Geir Gullestad Pettersen
 *
 */
public class BytecodeAgent implements ClassFileTransformer {

	private int ignored = 0;
	
	private int accounted = 0;
	
	private long bytesTotal = 0; 
	
	private static BytecodeAgent instance = null;
	
	private static boolean agentLoaded = false;
	
	private Map<String,byte[]> classDefinitions ;
	
	private PackageConfiguration config;
	
	private boolean firstTime = true;
	
	/**
	 * Private constructor
	 */
	private BytecodeAgent() {
		System.out.println( "BytecodeAgent constructor");
		classDefinitions = Collections.synchronizedMap( new HashMap<String,byte[]>() );		
		config = new PackageConfiguration();
	}
	
	/**
	 * Retrieves the instance 
	 * @return
	 */
	public synchronized static BytecodeAgent getInstance() {
		if( instance == null ) {
			instance = new BytecodeAgent();
		} 
		return instance;
	}
	
	/**
	 * Called by JVM to initialize agent.
	 * 
	 * @param agentArgs
	 * @param inst
	 */
	public static void premain(String agentArgs, Instrumentation inst) {
		
		System.out.println( "Initializing BytecodeAgent..");
		inst.addTransformer( getInstance() );		
		agentLoaded = true;
		System.out.println( "BytecodeAgent loaded" );
	}
	
	/* (non-Javadoc)
	 * @see java.lang.instrument.ClassFileTransformer#transform(java.lang.ClassLoader, java.lang.String, java.lang.Class, java.security.ProtectionDomain, byte[])
	 */
	public byte[] transform(ClassLoader loader, String className,
			Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
			byte[] classfileBuffer) throws IllegalClassFormatException {
		
		addByteCode(className, classfileBuffer);
		return null;
	}

	/**
	 * Returns <code>true</code> if <code>BytecodeAgent</code> is loaded into the JVM
	 * @return
	 */
	public static boolean isLoaded() {
		return agentLoaded;
	}
	
	/**
	 * Returns number (count) of classes that have been ignored.
	 * @return
	 */
	public int ignoredClasses() {
		return ignored;
	}
	
	public int accountedClasses()   {
		return accounted;
	}
	
	/**
     * Returns an array representing the bytecode for <code>className</code> if that
     * class has previously been loaded by a classloader.
     * 
	 * @param className
	 * @return
	 */
	byte[] getByteCode( String className ) {		
		if( firstTime ) {
			/* Get custom config the first time this method is being called
			 * Due to classloading issues this cannot be done in the constructor,
			 * transform or premain methods*/
			synchronized (classDefinitions) {
				PackageConfiguration customConfig = (PackageConfiguration) BeanFactory.getPackageConfiguration();
				/* replace default config if custom config was found */
				if( customConfig != null) {
					config = customConfig;
				}
				
				/* remove any classes that are not included in new configuration */
				filterDefinitions();
				
			}
			firstTime = false;
		}
		return classDefinitions.get( className );
	}
	
    
	private void filterDefinitions() {
		/* copy definitions to new map */
		Map<String,byte[]> filtered = Collections.synchronizedMap( new HashMap<String,byte[]>() );
		for (Entry<String,byte[]> entry : classDefinitions.entrySet()) {
			filtered.put( entry.getKey(), entry.getValue() );
		}				
		classDefinitions = filtered;
	}
	
	/**
     * Adds bytecode definition to the registry.
	 * @param className 
	 * @param byteCode
	 */
	private void addByteCode( String className, byte[] byteCode ) {
		synchronized( classDefinitions ) {
            // TODO use classloader to verify classname/bytecode
			if( classDefinitions.get( className ) != null ) {
				// definition already exist!
				System.err.println( className + " already in bytecode agent cache" );
			}
			else {
				if( config.includeClass( className )) { 
					classDefinitions.put( className, byteCode );
					accounted++;
					bytesTotal += byteCode.length;
				} else {
					ignored++;
				}
			}
		}
	}
}
